
from copy import deepcopy
from typing import List, Optional, Sequence, Union

import numpy as np
from mmengine.config import Config
from mmdet.structures import DetDataSample
from tqdm import tqdm
from ais_bench.infer.interface import InferSession

from ..postprocess import PostProcessor
from ..preprocess import ImagePreprocessor


class BaseInferencer:
    def __init__(self,
                 config: Union[str, Config] = None,
                 model: str = None,
                 device: Optional[int] = 0
                ) -> None:

        if not isinstance(config, Config):
            cfg = Config.fromfile(config)
        else:
            cfg = deepcopy(config)
        self.session = InferSession(device, model)
        self.batch_size = self.session.get_inputs()[0].shape[0]
        self.input_shape = self.session.get_inputs()[0].shape[2:]
        self.preprocessor = ImagePreprocessor(cfg, self.input_shape)
        self.postprocessor = PostProcessor(cfg)
        self.cfg = cfg

    def predict(self, batch_images) -> List[DetDataSample]:
        inputs, img_metas = self.preprocessor.preprocess_batch(
            batch_images, self.batch_size)
        outputs = self.session.infer(inputs)
        predictions = self.postprocessor.parse_outputs(
                            outputs, img_metas, return_datasample=True)
        return predictions
